#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
验证Milvus数据
"""

from pymilvus import connections, Collection, utility


def verify_data():
    """验证数据是否存在"""
    print("🔍 验证Docker Milvus中的数据")
    print("=" * 50)
    
    try:
        # 连接Milvus
        connections.connect("default", host="localhost", port="19530")
        print("✅ 成功连接到Milvus")
        
        # 列出所有集合
        collections = utility.list_collections()
        print(f"📚 找到 {len(collections)} 个集合: {collections}")
        
        # 检查knowledge_vectors集合
        if "knowledge_vectors" in collections:
            col = Collection("knowledge_vectors")
            print(f"\n🎯 knowledge_vectors集合:")
            print(f"  - 总记录数: {col.num_entities}")
            
            # 加载集合
            col.load()
            print("  - 集合已加载到内存")
            
            # 查询数据
            results = col.query(
                expr='filename like "C%"',  # 查找C开头的文件
                limit=3,
                output_fields=['filename', 'text']
            )
            
            print(f"  - 查询C开头的文件，找到 {len(results)} 条:")
            for result in results:
                print(f"    * {result['filename']}: {len(result['text'])} 字符")
            
            # 统计不同类型的文件
            all_results = col.query(
                expr='id != ""',
                output_fields=['filename']
            )
            
            file_types = {}
            for result in all_results:
                prefix = result['filename'][:2]
                file_types[prefix] = file_types.get(prefix, 0) + 1
            
            print(f"\n📊 文件类型统计:")
            for prefix, count in sorted(file_types.items()):
                print(f"  - {prefix}*: {count} 个文件")
            
            return True
        else:
            print("❌ 没有找到knowledge_vectors集合")
            return False
            
    except Exception as e:
        print(f"❌ 验证失败: {e}")
        return False


def search_demo():
    """演示搜索功能"""
    print("\n🔍 搜索演示")
    print("=" * 30)
    
    try:
        col = Collection("knowledge_vectors")
        col.load()
        
        # 搜索包含"模板"的文档
        results = col.query(
            expr='text like "%模板%"',
            limit=3,
            output_fields=['filename', 'text']
        )
        
        print(f"搜索包含'模板'的文档，找到 {len(results)} 条:")
        for i, result in enumerate(results, 1):
            print(f"\n{i}. 文件: {result['filename']}")
            text = result['text']
            # 找到"模板"关键词的位置
            pos = text.find("模板")
            if pos != -1:
                start = max(0, pos - 30)
                end = min(len(text), pos + 50)
                snippet = text[start:end]
                print(f"   内容: ...{snippet}...")
        
    except Exception as e:
        print(f"❌ 搜索失败: {e}")


if __name__ == "__main__":
    if verify_data():
        search_demo()
    
    print("\n" + "=" * 50)
    print("✅ 数据验证完成")
    print("您的知识库数据确实存储在Docker Milvus中！")
    print("如果您在其他工具中看不到数据，可能是:")
    print("1. 连接的是错误的Milvus实例")
    print("2. 查看的是错误的集合名称")
    print("3. 需要先加载集合到内存")
    print("4. 使用了错误的查询语法")
